package huffmancode;

import javax.swing.*;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.io.*;
import java.util.*;
import java.util.List;
import java.util.Queue;

/**
 * @description:
 * @author: 陈瑞东
 * @create: 2021-08-13 11:29
 */
public class HuffmanCode {
    public static void main(String[] args) throws Exception {
        start();
    }

    public static void start() throws Exception {
        UIManager.setLookAndFeel("com.sun.java.swing.plaf.windows.WindowsLookAndFeel");

        JFrame jFrame = new JFrame("赫夫曼编码");

        Box hBox = Box.createHorizontalBox();

        Box vBox = Box.createVerticalBox();

        Box firstLine = Box.createHorizontalBox();
        JLabel title = new JLabel("赫夫曼编码");
        title.setFont(new Font("Stsong", Font.BOLD, 30));
        firstLine.add(title);

        Box secondInsetLine = Box.createHorizontalBox();
        JLabel inputTip = new JLabel("解压时输入文件名");
        inputTip.setFont(new Font("Stsong", Font.PLAIN, 24));
        secondInsetLine.add(inputTip);

        Box secondLine = Box.createHorizontalBox();
        JTextField fileNameField = new JTextField(20);
        JFileChooser fileChooser = new JFileChooser(".");
        secondLine.add(fileNameField);

        Box ThirdLine = Box.createHorizontalBox();
        JButton zipBtn = new JButton("压缩");
        JButton unZipBtn = new JButton("解压");
        zipBtn.setFont(new Font("Stsong", Font.PLAIN, 24));
        unZipBtn.setFont(new Font("Stsong", Font.PLAIN, 24));

        ThirdLine.add(zipBtn);
        ThirdLine.add(Box.createHorizontalStrut(20));
        ThirdLine.add(unZipBtn);

        vBox.add(Box.createVerticalStrut(20));
        vBox.add(firstLine);
        vBox.add(Box.createVerticalStrut(20));
        vBox.add(secondInsetLine);
        vBox.add(Box.createVerticalStrut(20));
        vBox.add(secondLine);
        vBox.add(Box.createVerticalStrut(20));
        vBox.add(ThirdLine);
        vBox.add(Box.createVerticalStrut(20));

        hBox.add(Box.createHorizontalStrut(20));
        hBox.add(vBox);
        hBox.add(Box.createHorizontalStrut(20));


        jFrame.add(hBox);

        zipBtn.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                int result = fileChooser.showOpenDialog(jFrame);
                if (result == JFileChooser.CANCEL_OPTION) {
                    return;
                }
                File file = fileChooser.getSelectedFile();

                if (file == null) {
                    return;
                }
                String text = file.getAbsolutePath();
                if (!file.exists() || file.isDirectory()) {
                    JOptionPane.showMessageDialog(jFrame, "请选择一个文件！", "提示", JOptionPane.WARNING_MESSAGE);
                    return;
                }
                zipFile(text, text + ".zip");
                JOptionPane.showMessageDialog(jFrame, "压缩成功！", "提示", JOptionPane.INFORMATION_MESSAGE);
            }
        });

        unZipBtn.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                int result = fileChooser.showOpenDialog(jFrame);
                if (result == JFileChooser.CANCEL_OPTION) {
                    return;
                }
                File file = fileChooser.getSelectedFile();
                if (file == null) {
                    return;
                }
                if (!file.exists() || file.isDirectory()) {
                    JOptionPane.showMessageDialog(jFrame, "请选择一个文件！", "提示", JOptionPane.WARNING_MESSAGE);
                }

                String filename = fileNameField.getText();
                if (filename == null || "".equals(filename)) {
                    JOptionPane.showMessageDialog(jFrame, "请输入解压出的文件名！", "提示", JOptionPane.WARNING_MESSAGE);
                    return;
                }

                File destFile = new File(file.getParentFile(), filename);

                unZipFile(file.getAbsolutePath(), destFile.getAbsolutePath());
                JOptionPane.showMessageDialog(jFrame, "解压成功！", "提示", JOptionPane.INFORMATION_MESSAGE);
            }
        });


        jFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
//        jFrame.pack();
        jFrame.setSize(400, 300);
        jFrame.setResizable(false);
        jFrame.setLocationRelativeTo(null);
        jFrame.setVisible(true);
    }

    public static void unZipFile(String src, String dest) {
        try (
                FileInputStream in = new FileInputStream(src);
                ObjectInputStream ois = new ObjectInputStream(in);
                FileOutputStream out = new FileOutputStream(dest);

        ){
            byte[] bytes = (byte[]) ois.readObject();
            Map<Byte, String> huffmanCodes = (Map<Byte, String>) ois.readObject();
            byte[] outBytes = decode(bytes, huffmanCodes);

            out.write(outBytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void zipFile(String src, String dest) {
        try (
            FileInputStream in = new FileInputStream(src);
            FileOutputStream out = new FileOutputStream(dest);
            ObjectOutputStream oos = new ObjectOutputStream(out);

        ){
            byte[] data = new byte[in.available()];
            in.read(data);
            Map<Byte, String> huffmanCodes = getCodes(data);
            byte[] outBytes = zip(data, huffmanCodes);

            oos.writeObject(outBytes);
            oos.writeObject(huffmanCodes);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static byte[] decode(byte[] huffmanCodesBytes, Map<Byte, String> huffmanCodes) {
        List<Byte> ansList = new ArrayList<>();
        Map<String, Byte> reverseMap = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            reverseMap.put(entry.getValue(), entry.getKey());
        }

        StringBuilder builder = new StringBuilder();
        for (int i = 1; i < huffmanCodesBytes.length; i++) {
            byte code = huffmanCodesBytes[i];
            int offset = 8;
            if (i == huffmanCodesBytes.length - 1) {
                offset = huffmanCodesBytes[0];
            }
            String str = Integer.toBinaryString(1 << offset | code);
            str = str.substring(str.length() - offset);
            for (int j = 0; j < offset; j++) {
                builder.append(str.charAt(j));
                if (reverseMap.containsKey(builder.toString())) {
                    ansList.add(reverseMap.get(builder.toString()));
                    builder = new StringBuilder();
                }
            }
        }

        byte[] ans = new byte[ansList.size()];
        for (int i = 0; i < ans.length; i++) {
            ans[i] = ansList.get(i);
        }
        return ans;
    }

    public static byte[] zip(byte[] contentBytes, Map<Byte, String> huffmanCodes) {
        List<Byte> ansList = new ArrayList<>();
        StringBuilder builder = new StringBuilder();

        for (byte contentByte : contentBytes) {
            builder.append(huffmanCodes.get(contentByte));
            if (builder.length() >= 8) {
                ansList.add((byte) Integer.parseInt(builder.substring(0, 8), 2));
                builder.delete(0, 8);
            }
        }
        while (builder.length() >= 8) {
            ansList.add((byte) Integer.parseInt(builder.substring(0, 8), 2));
            builder.delete(0, 8);
        }
        if (builder.length() != 0) {
            ansList.add((byte) Integer.parseInt(builder.toString(), 2));
            ansList.add(0, (byte) builder.length());
        } else {
            ansList.add(0, (byte) 8);
        }
        byte[] ans = new byte[ansList.size()];
        for (int i = 0; i < ans.length; i++) {
            ans[i] = ansList.get(i);
        }

        return ans;
    }

    public static Map<Byte, String> getCodes(byte[] contentBytes) {
        List<Node> nodes = getNodes(contentBytes);
        Node huffmanTree = creatHuffmanTree(nodes);
        Map<Byte, String> huffmanCodes = getCodes(huffmanTree);
        return huffmanCodes;
    }

    public static Map<Byte, String> getCodes(Node huffmanTree) {
        Map<Byte, String> ans = new HashMap<>();
        Queue<Object[]> queue = new LinkedList<>();
        queue.offer(new Object[] {"", huffmanTree});
        while (!queue.isEmpty()) {
            Object[] data = queue.poll();
            String code = (String) data[0];
            Node node = (Node) data[1];
            if (node.getLeft() != null) {
                queue.offer(new Object[] {code + "0", node.getLeft()});
            }
            if (node.getRight() != null) {
                queue.offer(new Object[] {code + "1", node.getRight()});
            }
            if (node.getLeft() == null && node.getRight() == null) {
                ans.put(node.getData(), code);
            }
        }
        return ans;
    }

    public static Node creatHuffmanTree(List<Node> nodes) {
        while (nodes.size() != 1) {
            Collections.sort(nodes);
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            Node parentNode = new Node(leftNode.getWeight() + rightNode.getWeight(), null);
            parentNode.setLeft(leftNode);
            parentNode.setRight(rightNode);
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            nodes.add(parentNode);
        }
        return nodes.get(0);
    }

    public static List<Node> getNodes(byte[] bytes) {
        Map<Byte, Integer> counts = new HashMap<>();
        List<Node> nodes = new ArrayList<>();
        for (byte b : bytes) {
            counts.put(b, counts.getOrDefault(b, 0) + 1);
        }

        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new Node(entry.getValue(), entry.getKey()));
        }
        return nodes;
    }
}

class Node implements Comparable<Node>{
    private int weight;
    private Byte data;
    private Node left;
    private Node right;

    public Node(int weight, Byte data) {
        this.weight = weight;
        this.data = data;
    }

    public int getWeight() {
        return weight;
    }

    public void setWeight(int weight) {
        this.weight = weight;
    }

    public Byte getData() {
        return data;
    }

    public void setData(Byte data) {
        this.data = data;
    }

    public Node getLeft() {
        return left;
    }

    public void setLeft(Node left) {
        this.left = left;
    }

    public Node getRight() {
        return right;
    }

    public void setRight(Node right) {
        this.right = right;
    }

    public String toString() {
        return "Node [data = " + data + " weight=" + weight + "]";
    }

    //前序遍历
    public void preOrder() {
        System.out.println(this);
        if(this.left != null) {
            this.left.preOrder();
        }
        if(this.right != null) {
            this.right.preOrder();
        }
    }
    @Override
    public int compareTo(Node o) {
        return this.weight - o.weight;
    }
}



